# analyzer.py

import cv2
import os
import time
from datetime import datetime
from ultralytics import YOLO
from config import ANALYSIS_DURATION, SAMPLE_FPS, YOLO_MODEL_PATH
from utils.minio_client import upload_file
from utils.qwen_client import analyze_keyframe_with_qwen

# def upload_file(local_path: str, object_name: str) -> str:
#     """
#     上传文件到 MinIO
#     """
#     return f"s3://www.ai.construction/{object_name}"
  
    
# 加载 YOLO 模型
CONSTRUCTION_CLASSES = {
        'worker',
        'hardhat',
        'excavator',
        'barrier',
        'scaffolding',
        "Dump truck",
        "Excavator",
        "Motor grader",
        "Roller",
        "Crane manipulator",
        "Gazelle",
        "Forklift Standart",
        "Bucket loader Big",
        "Mixer",
        "Tanker",
        "Bulldozer",
        "Cleaning equipment",
        "Truck",
        "Trailer",
        "Forklift Giraffe",
        "Bucket loader Standart",
        "Autocran",
        "Truck crane"
    }

# 类别颜色（BGR）
COLORS = {
    'worker': (0, 255, 255),      # 黄色
    'hardhat': (0, 255, 0),       # 绿色
    'excavator': (255, 0, 0),     # 红色
    'barrier': (255, 165, 0),     # 橙色
    'scaffolding': (128, 0, 128)  # 紫色
}

model = YOLO(YOLO_MODEL_PATH)
def analyze_single_camera(camera_id: str, rtsp_url: str,device: str = 'cuda'):
    """
    分析单路摄像头：30秒视频 → YOLO → Qwen-VL → 仅施工事件上传 MinIO
    """
    model.to(device)
    
    t = int(time.time())
    temp_dir = f"/tmp/analysis_{camera_id}"
    os.makedirs(temp_dir, exist_ok=True)

    video_file = os.path.join(temp_dir, f"{t}_detection_.mp4")
    keyframe_file = os.path.join(temp_dir, f"{t}_keyframe.jpg")

    cap = cv2.VideoCapture(rtsp_url)
    cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
    #cap.set(cv2.CAP_PROP_READ_TIMEOUT, 10000)

    if not cap.isOpened():
        return {"camera_id": camera_id, "error": "无法拉流", "uploaded": False}

    fps = max(cap.get(cv2.CAP_PROP_FPS), 15)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(video_file, fourcc, fps, (width, height))

    start_time = time.time()
    frame_count = 0
    check_frames = 0
    construction_frames = 0
    best_frame = None
    best_score = 0
    best_detections = []

    while (time.time() - start_time) < ANALYSIS_DURATION:
        ret, frame = cap.read()
        if not ret:
            break

        frame_count += 1
        detections = []
        frame_area = width * height

        # if frame_count % int(fps / SAMPLE_FPS) != 0:
        #     continue
        
        check_frames += 1
        
        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        resized = cv2.resize(rgb_frame, (640, 640))

        results = model(resized, imgsz=640, conf=0.5, device=device,  verbose=False)
        annotated = results[0].plot()
        annotated = cv2.cvtColor(annotated, cv2.COLOR_RGB2BGR)
        
         # 场景判断
        boxes = results[0].boxes.cpu().numpy() #results.boxes.cpu().numpy() 
        has_construction = any(
            results[0].names[int(box.cls[0])] in CONSTRUCTION_CLASSES for box in boxes
        )
        #print(f"模型结果: {has_construction}'\r'")
        if has_construction:
            construction_frames += 1
            
            # 计算评分，选择最佳关键帧
            
            conf_sum = sum(float(box.conf[0]) for box in boxes)
            score = conf_sum * 1.5 + len(boxes) * 2.0
            if score > best_score:
                best_score = score
                best_frame = annotated.copy()
                best_detections = [
                    {'class': results[0].names[int(box.cls[0])], 'conf': float(box.conf[0])}
                    for box in boxes
                ]
            
            #print(f"最佳:{best_score},施工比例: {score}, 检查帧: {check_frames} ，施工帧: {construction_frames}，总置信度: {conf_sum}，目标数: {len(boxes)}")

            construction_frames += 1
            cv2.putText(annotated, "CONSTRUCTION SITE DETECTED!", (50, 50),
                        cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0, 0, 255), 3)
        else:
            cv2.putText(annotated, "NA", (10, height - 30),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2, cv2.LINE_AA)
        
        # 写入视频
        #out.write(frame)
        out.write(cv2.resize(annotated, (width, height)))
        # 实时显示（可选）
        #cv2.imshow('Construction Detection', cv2.resize(annotated_bgr, (640, 640)))
        # if cv2.waitKey(1) & 0xFF == ord('q'):
        #     break
        
        #print(f"处理进度: {frame_count}/10000 帧 {construction_frames} ，检查帧: {check_frames}", end='\n')

    cap.release()
    out.release()

    ratio = construction_frames #/ max(check_frames, 1)
    print(f"施工比例: {ratio} ，检查帧: {check_frames} ，施工帧: {construction_frames}")
    is_suspicious = ratio > 0.2

    if not is_suspicious or best_frame is None:
        print(f"非施工退出: {ratio}")
        cleanup_temp(temp_dir, [video_file, keyframe_file])
        return {
            "camera_id": camera_id,
            "construction_ratio": ratio,
            "is_construction": False,
            "uploaded": False
        }

    cv2.imwrite(keyframe_file, cv2.resize(best_frame, (width, height)))

    qwen_result = analyze_keyframe_with_qwen(
        keyframe_file,
        {
            "ratio": ratio,
            "count": construction_frames,
            "detections": best_detections
        }
    )

    is_construction = qwen_result.get("is_construction", False)
    upload_result = {"video_": None, "keyframe": None}

    if is_construction:
        date_str = datetime.now().strftime("%Y-%m-%d")
        video_obj = f"{date_str}/videos/{camera_id}/detection_{t}.mp4"
        keyframe_obj = f"{date_str}/keyframes/{camera_id}/keyframe_{t}.jpg"

        upload_result["video"] = upload_file(video_file, video_obj)
        upload_result["keyframe"] = [ upload_file(keyframe_file, keyframe_obj) ]

    cleanup_temp(temp_dir, [video_file, keyframe_file])

    return {
        "camera_id": camera_id,
        "is_construction": is_construction,
        "construction_ratio": ratio,
        "qwen_reason": qwen_result.get("reason", ""),
        "uploaded": is_construction,
        "uploads": upload_result
    }

def cleanup_temp(temp_dir: str, files: list):
    return
    for f in files:
        if os.path.exists(f):
            os.remove(f)
    if os.path.exists(temp_dir):
        os.rmdir(temp_dir)